import chatbot.config as config
import fasttext

class Classify:
    def __init__(self):
        """
        加载训练好的model
        """
        self.model = fasttext.load_model(config.classify_model_final_path)
        self.model_by_word = fasttext.load_model(config.classify_model_path_by_word)

    def predict(self, sentence):
        """
        预测输入数据的结果(是不是qa)，准确率
        :param sentence:{"cut":str,"cut_by_word":str}
        :return: (label,acc)
        """
        result = self.model.predict(sentence["cut"])
        result_by_word = self.model_by_word.predict(sentence["cut_by_word"])
        # *result解包，result里是两个返回值
        for label, acc, label_by_word, acc_by_word in zip(*result, *result_by_word):
            # 两个类别，只适用于当前demo
            if label == "__label__chat":
                label = "__label__QA"
                acc = 1 - acc
            if label_by_word == "__label__chat":
                label_by_word = "__label__QA"
                acc_by_word = 1 - acc_by_word
            # 判断准确率
            if acc > 0.95 or acc_by_word > 0.95:
                return "QA", max(acc, acc_by_word)
            else:
                return "chat", 1 - min(acc, acc_by_word)

            # 三个以上类别
            # if label == label_by_word:
            #     if acc > 0.95 or acc_by_word > 0.95:
            #         return label, max(acc, acc_by_word)
            #     else:
            #         return None, 0  # 结果不符合
            # else:
            #     if acc_by_word > 0.99:
            #         return label_by_word, acc_by_word
            #     elif acc > 0.98:
            #         return label, acc
            #     else:
            #         return None, 0
